"""pointnet2 test"""

from mind3d.utils.pointnet2_util import farthest_point_sample, square_distance, index_points
from mindspore import Tensor
import mindspore as ms
import numpy as np
import matplotlib.pylab as mp


def test_square_distance():
    """
    test the function of square_distance
    """
    src = Tensor(np.array([[[0, 0, 0], [3, 4, 0]]])).astype(float)
    dst = Tensor(np.array([[[3, 4, 0], [0, 0, 0]]])).astype(float)
    truth = Tensor(np.array([[[5, 0], [0, 5]]])).astype(float)
    result = square_distance(src, dst)
    assert result == truth


def test_index_points():
    """
    test the function of index_points
    """
    points = Tensor(np.arange(1, 101)[:60].reshape((2, 10, 3)), ms.float32)
    idx = Tensor(np.arange(1, 20)[:6].reshape((2, 3)), ms.int32)
    new_points = index_points(points, idx)
    assert new_points.shape == [2, 3 ,3]


def test_farthest_point_sample(npoint=100):
    """
    test the function of farthest_point_sample
    """
    # x, y, z = np.mgrid[100:250:15, 200:350:15, 400:600:10]
    # xyz = np.array([x.reshape(1, -1), y.reshape(1, -1), z.reshape(1, -1)])
    # xyz = xyz.transpose(1, 2, 0)
    # print(xyz)
    # # draw
    # mp.figure("3D Scatter", facecolor="lightgray")
    # ax3d = mp.gca(projection="3d")
    # mp.title('3D FPS', fontsize=20)
    # ax3d.set_xlabel('x', fontsize=14)
    # ax3d.set_ylabel('y', fontsize=14)
    # ax3d.set_zlabel('z', fontsize=14)
    # mp.tick_params(labelsize=10)
    # ax3d.scatter(xyz[0, :, 0], xyz[0, :, 1], xyz[0, :, 2], cmap="jet", marker="o")
    #
    # sim_data = Tensor(xyz).astype(np.float32)
    # centroids = farthest_point_sample(xyz=sim_data, npoint=npoint)
    # print("Sampled pts: ", centroids)
    # centroids = centroids.asnumpy().astype(int)
    # for i in range(centroids.shape[1]):
    #     print(xyz[0, centroids[0, i], :])
    #     ax3d.scatter(xyz[0, centroids[0, i], 0], xyz[0, centroids[0, i], 1], xyz[0, centroids[0, i], 2], c='r')
    # mp.show()
    points = Tensor(np.arange(1, 154500)[:150000].reshape((10, 5000, 3)), ms.float32)
    npoint = 1024
    new_points = farthest_point_sample(points, npoint)
    assert new_points.shape == [10, 1024 ,3]



if __name__ == '__main__':
    test_square_distance()
    test_farthest_point_sample()
    test_index_points()
